import cv2
import os
import sys
import numpy as np
from tqdm import tqdm
import tkinter as tk
from tkinter import filedialog

# 设置控制台输出编码为utf-8
sys.stdout.reconfigure(encoding='utf-8')

def remove_video_watermark(input_video_path, output_video_path, watermark_height=50, watermark_width=None, position="bottom_right", method="inpaint", inpaint_radius=10, inpaint_method=cv2.INPAINT_NS):
    """
    去除视频右下角的水印
    
    参数:
        input_video_path: 输入视频路径
        output_video_path: 输出视频路径
        watermark_height: 水印的高度（默认50像素）
        watermark_width: 水印的宽度（默认为None，表示使用视频宽度的1/4）
        position: 水印位置，可选 "bottom_right", "bottom_left", "top_right", "top_left"
        method: 去水印方法，可选 "inpaint"(图像修复), "blur"(模糊处理), "color"(颜色填充)
        inpaint_radius: 图像修复算法的半径参数，值越大修复范围越广（默认10）
        inpaint_method: 图像修复算法，可选 cv2.INPAINT_NS（基于Navier-Stokes的方法）或 cv2.INPAINT_TELEA（基于快速行进算法）
    """
    try:
        # 打开视频文件
        cap = cv2.VideoCapture(input_video_path)
        if not cap.isOpened():
            print(f"错误：无法打开视频文件 {input_video_path}")
            return False
        
        # 获取视频的基本信息
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        fps = cap.get(cv2.CAP_PROP_FPS)
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        
        # 如果未指定水印宽度，则默认使用视频宽度的1/4
        if watermark_width is None:
            watermark_width = width // 4
        
        # 创建视频写入器
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # 使用MP4编码
        out = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))
        
        print(f"开始处理视频，总帧数: {total_frames}")
        print(f"使用 {method} 方法去除水印")
        
        # 使用tqdm显示进度条
        for _ in tqdm(range(total_frames)):
            ret, frame = cap.read()
            if not ret:
                break
                
            # 确定水印区域坐标
            if position == "bottom_right":
                # 右下角
                x1, y1 = width - watermark_width, height - watermark_height
                x2, y2 = width, height
            elif position == "bottom_left":
                # 左下角
                x1, y1 = 0, height - watermark_height
                x2, y2 = watermark_width, height
            elif position == "top_right":
                # 右上角
                x1, y1 = width - watermark_width, 0
                x2, y2 = width, watermark_height
            elif position == "top_left":
                # 左上角
                x1, y1 = 0, 0
                x2, y2 = watermark_width, watermark_height
            
            # 根据选择的方法去除水印
            if method == "inpaint":
                # 方法1：图像修复 - 使用OpenCV的inpaint函数
                
                # 创建水印区域的遮罩
                mask = np.zeros(frame.shape[:2], dtype=np.uint8)
                
                # 增强型遮罩处理：使用边缘检测来更精确地识别水印区域
                roi = frame[y1:y2, x1:x2].copy()
                gray_roi = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
                
                # 使用阈值分割找出可能的水印区域
                _, thresh = cv2.threshold(gray_roi, 240, 255, cv2.THRESH_BINARY)
                
                # 将阈值结果与原始遮罩结合
                mask[y1:y2, x1:x2] = 255  # 首先将整个区域标记为需要修复
                
                # 使用选定的图像修复算法进行修复
                frame = cv2.inpaint(frame, mask, inpaint_radius, inpaint_method)
                
            elif method == "blur":
                # 方法2：模糊处理 - 对水印区域进行高斯模糊
                # 提取水印区域
                watermark_region = frame[y1:y2, x1:x2].copy()
                
                # 对水印区域进行高斯模糊
                blurred = cv2.GaussianBlur(watermark_region, (25, 25), 0)
                
                # 将模糊后的区域放回原图
                frame[y1:y2, x1:x2] = blurred
                
            elif method == "color":
                # 方法3：颜色填充 - 使用周围区域的平均颜色填充
                # 获取水印区域附近的平均颜色
                if y1 > 10:
                    # 如果水印在底部，则取上方区域的平均颜色
                    avg_color = np.mean(frame[y1-10:y1, x1:x2], axis=(0, 1)).astype(np.uint8)
                else:
                    # 如果水印在顶部，则取下方区域的平均颜色
                    avg_color = np.mean(frame[y2:min(y2+10, height), x1:x2], axis=(0, 1)).astype(np.uint8)
                
                # 用平均颜色填充水印区域
                frame[y1:y2, x1:x2] = avg_color
            
            # 写入处理后的帧
            out.write(frame)
        
        # 释放资源
        cap.release()
        out.release()
        
        print(f"视频处理完成，已保存为: {output_video_path}")
        return True
        
    except Exception as e:
        print(f"处理视频时发生错误: {str(e)}")
        return False

def analyze_watermark(input_video_path, frame_number=0, watermark_height=50, watermark_width=None, position="bottom_right"):
    """
    分析视频中的水印区域，并显示一帧以便调整参数
    
    参数:
        input_video_path: 输入视频路径
        frame_number: 要分析的帧号
        watermark_height: 水印的高度
        watermark_width: 水印的宽度
        position: 水印位置
    """
    try:
        # 打开视频文件
        cap = cv2.VideoCapture(input_video_path)
        if not cap.isOpened():
            print(f"错误：无法打开视频文件 {input_video_path}")
            return False
        
        # 获取视频的基本信息
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        
        # 如果未指定水印宽度，则默认使用视频宽度的1/4
        if watermark_width is None:
            watermark_width = width // 4
        
        # 跳转到指定帧
        cap.set(cv2.CAP_PROP_POS_FRAMES, frame_number)
        ret, frame = cap.read()
        if not ret:
            print(f"错误：无法读取第 {frame_number} 帧")
            return False
        
        # 确定水印区域坐标
        if position == "bottom_right":
            x1, y1 = width - watermark_width, height - watermark_height
            x2, y2 = width, height
        elif position == "bottom_left":
            x1, y1 = 0, height - watermark_height
            x2, y2 = watermark_width, height
        elif position == "top_right":
            x1, y1 = width - watermark_width, 0
            x2, y2 = width, watermark_height
        elif position == "top_left":
            x1, y1 = 0, 0
            x2, y2 = watermark_width, watermark_height
        
        # 在原始帧上绘制水印区域的矩形框
        frame_with_rect = frame.copy()
        cv2.rectangle(frame_with_rect, (x1, y1), (x2, y2), (0, 0, 255), 2)
        
        # 获取输入视频的目录
        input_dir = os.path.dirname(input_video_path)
        if not input_dir:
            input_dir = os.getcwd()
        
        # 保存带有矩形框的帧
        output_path = os.path.join(input_dir, "watermark_analysis.jpg")
        cv2.imwrite(output_path, frame_with_rect)
        
        # 释放资源
        cap.release()
        
        print(f"水印分析完成，已保存为: {output_path}")
        print(f"水印区域坐标: 左上角({x1}, {y1}), 右下角({x2}, {y2})")
        print(f"水印大小: 宽度={watermark_width}, 高度={watermark_height}")
        
        return True
        
    except Exception as e:
        print(f"分析水印时发生错误: {str(e)}")
        return False

def select_video_file():
    """
    打开文件选择对话框，让用户选择要处理的视频文件
    
    返回:
        选择的文件路径，如果用户取消则返回None
    """
    # 创建一个临时的Tkinter根窗口
    root = tk.Tk()
    root.withdraw()  # 隐藏主窗口
    
    # 打开文件选择对话框
    file_path = filedialog.askopenfilename(
        title="选择要去除水印的视频文件",
        filetypes=[
            ("视频文件", "*.mp4 *.avi *.mov *.mkv *.wmv"),
            ("所有文件", "*.*")
        ]
    )
    
    # 销毁临时窗口
    root.destroy()
    
    return file_path if file_path else None

def generate_output_path(input_path):
    """
    根据输入文件路径生成输出文件路径，在文件名后添加"end"后缀
    
    参数:
        input_path: 输入文件路径
        
    返回:
        输出文件路径
    """
    # 获取文件目录、文件名和扩展名
    directory = os.path.dirname(input_path)
    filename = os.path.basename(input_path)
    name, ext = os.path.splitext(filename)
    
    # 创建新的文件名
    new_filename = f"{name}end{ext}"
    
    # 如果目录为空，使用当前目录
    if not directory:
        directory = os.getcwd()
    
    # 组合新的文件路径
    output_path = os.path.join(directory, new_filename)
    
    return output_path

if __name__ == "__main__":
    print("视频去水印工具")
    print("==============")
    
    # 让用户选择要处理的视频文件
    print("请选择要处理的视频文件...")
    input_video = select_video_file()
    
    if not input_video:
        print("未选择任何文件，程序退出。")
        sys.exit(0)
    
    # 生成输出文件路径
    output_video = generate_output_path(input_video)
    
    print(f"已选择视频: {input_video}")
    print(f"处理后的视频将保存为: {output_video}")
    
    # 检查输入文件是否存在
    if not os.path.exists(input_video):
        print(f"错误：找不到输入文件 {input_video}")
        sys.exit(1)
    
    # 首先分析水印位置和大小
    print("\n正在分析水印...")
    analyze_watermark(
        input_video, 
        frame_number=0,  # 分析第一帧
        watermark_height=50,  # 水印高度，可以根据实际情况调整
        watermark_width=200,  # 水印宽度，可以根据实际情况调整
        position="bottom_right"  # 水印位置
    )
    
    # 去除水印
    print("\n正在去除水印...")
    success = remove_video_watermark(
        input_video, 
        output_video, 
        watermark_height=50,  # 水印高度，可以根据实际情况调整
        watermark_width=200,  # 水印宽度，可以根据实际情况调整
        position="bottom_right",  # 水印位置
        method="inpaint",  # 使用图像修复方法
        inpaint_radius=10,  # 修复半径
        inpaint_method=cv2.INPAINT_NS  # 使用Navier-Stokes方法，通常效果更好
    )
    
    if success:
        print(f"视频水印去除成功！输出文件: {output_video}")
    else:
        print("视频水印去除失败，请检查错误信息。")